// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/executors/io_thread_pool_executor.h>
#include <gtest/gtest.h>

#include <pollux/common/caching/ssd_cache.h>
#include <pollux/core/expressions.h>
#include <pollux/core/plan_node.h>
#include <pollux/exec/hash_probe.h>
#include <pollux/testing/exec/util/query_assertions.h>
#include <pollux/parse/expressions_parser.h>
#include <pollux/vector/vector_maker.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::exec::test {
class OperatorTestBase : public virtual testing::Test,
                         public pollux::VectorBuilder {
 public:
  /// The following methods are used by google unit test framework to do
  /// one-time setup/teardown for all the unit tests from OperatorTestBase. We
  /// make them public as some benchmark like ReduceAgg also call these methods
  /// to setup/teardown benchmark test environment.
  static void SetUpTestCase();
  static void TearDownTestCase();

  /// Sets up the pollux memory system.
  ///
  /// NOTE: a second call to this will clear the previous memory system
  /// instances and create a new set.
  static void setupMemory(
      int64_t allocatorCapacity,
      int64_t arbitratorCapacity,
      int64_t arbitratorReservedCapacity,
      int64_t memoryPoolInitCapacity,
      int64_t memoryPoolReservedCapacity,
      int64_t memoryPoolMinReclaimBytes,
      int64_t memoryPoolAbortCapacityLimit);

  static void resetMemory();

 protected:
  OperatorTestBase();
  ~OperatorTestBase() override;

  void SetUp() override;

  void TearDown() override;

  /// Allow base classes to register custom vector serde.
  /// By default, registers Presto-compatible serde.
  virtual void registerVectorSerde();

  void createDuckDbTable(const std::vector<RowVectorPtr>& data) {
    duckDbQueryRunner_.createTable("tmp", data);
  }

  void createDuckDbTable(
      const std::string& tableName,
      const std::vector<RowVectorPtr>& data) {
    duckDbQueryRunner_.createTable(tableName, data);
  }

  std::shared_ptr<Task> assertQueryOrdered(
      const core::PlanNodePtr& plan,
      const std::string& duckDbSql,
      const std::vector<uint32_t>& sortingKeys) {
    return test::assertQuery(plan, duckDbSql, duckDbQueryRunner_, sortingKeys);
  }

  std::shared_ptr<Task> assertQueryOrdered(
      const CursorParameters& params,
      const std::string& duckDbSql,
      const std::vector<uint32_t>& sortingKeys) {
    return test::assertQuery(
        params, [&](auto*) {}, duckDbSql, duckDbQueryRunner_, sortingKeys);
  }

  /// Assumes plan has a single leaf node. All splits are added to that node.
  std::shared_ptr<Task> assertQueryOrdered(
      const core::PlanNodePtr& plan,
      const std::vector<std::shared_ptr<connector::ConnectorSplit>>& splits,
      const std::string& duckDbSql,
      const std::vector<uint32_t>& sortingKeys) {
    return assertQuery(plan, splits, duckDbSql, sortingKeys);
  }

  std::shared_ptr<Task> assertQuery(
      const CursorParameters& params,
      const std::string& duckDbSql) {
    return test::assertQuery(
        params, [&](exec::Task* /*task*/) {}, duckDbSql, duckDbQueryRunner_);
  }

  std::shared_ptr<Task> assertQuery(
      const core::PlanNodePtr& plan,
      const std::string& duckDbSql) {
    return test::assertQuery(plan, duckDbSql, duckDbQueryRunner_);
  }

  std::shared_ptr<Task> assertQuery(
      const core::PlanNodePtr& plan,
      const RowVectorPtr& expectedResults) {
    return test::assertQuery(plan, {expectedResults});
  }

  std::shared_ptr<Task> assertQuery(
      const CursorParameters& params,
      const RowVectorPtr& expectedResults) {
    return test::assertQuery(params, {expectedResults});
  }

  /// Assumes plan has a single leaf node. All splits are added to that node.
  std::shared_ptr<Task> assertQuery(
      const core::PlanNodePtr& plan,
      const std::vector<std::shared_ptr<connector::ConnectorSplit>>&
          connectorSplits,
      const std::string& duckDbSql,
      std::optional<std::vector<uint32_t>> sortingKeys = std::nullopt);

  /// Assumes plan has a single leaf node. All splits are added to that node.
  std::shared_ptr<Task> assertQuery(
      const core::PlanNodePtr& plan,
      std::vector<exec::Split>&& splits,
      const std::string& duckDbSql,
      std::optional<std::vector<uint32_t>> sortingKeys = std::nullopt);

  std::shared_ptr<Task> assertQuery(
      const core::PlanNodePtr& plan,
      std::unordered_map<core::PlanNodeId, std::vector<exec::Split>>&& splits,
      const std::string& duckDbSql,
      std::optional<std::vector<uint32_t>> sortingKeys = std::nullopt);

  static RowTypePtr make_row_type(std::vector<TypePtr>&& types) {
    return pollux::VectorMaker::row_type(
        std::forward<std::vector<TypePtr>&&>(types));
  }

  static std::shared_ptr<core::FieldAccessTypedExpr> toFieldExpr(
      const std::string& name,
      const RowTypePtr& rowType);

  core::TypedExprPtr parseExpr(
      const std::string& text,
      RowTypePtr rowType,
      const parse::ParseOptions& options = {});

 public:
  static void deleteTaskAndCheckSpillDirectory(std::shared_ptr<Task>& task);

 protected:
  // Used as default AsyncDataCache.
  static inline std::shared_ptr<cache::AsyncDataCache> asyncDataCache_;

  DuckDbQueryRunner duckDbQueryRunner_;

  // Used for driver thread execution.
  std::unique_ptr<melon::CPUThreadPoolExecutor> driverExecutor_;

  // Used for IO prefetch and spilling.
  std::unique_ptr<melon::IOThreadPoolExecutor> ioExecutor_;
};
} // namespace kumo::pollux::exec::test
